package com.mxx.common.thread;

import com.mxx.common.core.ResetCounter;

import java.util.ArrayList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.Consumer;

/**
 * @author: mxx
 * @date: 2021/8/19
 * @time: 11:21
 * @Description: 数据批处理队列
 */
public class BatchProcessQueue<T> {
    private BlockingQueue<T> batchQueue;
    private ResetCounter counter;
    private Integer batchNumber;
    private Consumer<List<T>> batchExecuteHandler;

    public BatchProcessQueue(BlockingQueue<T> batchQueue, Integer batchNumber, Consumer<List<T>> batchExecuteHandler) {
        this.batchQueue = batchQueue;
        this.batchNumber = batchNumber;
        this.batchExecuteHandler = batchExecuteHandler;
        this.counter = new ResetCounter(batchNumber);
    }

    public BatchProcessQueue(Integer batchNumber, Consumer<List<T>> batchExecuteHandler) {
        this(new LinkedBlockingQueue(batchNumber*3), batchNumber, batchExecuteHandler);
    }

    public void add(T t) {
        if (batchQueue.add(t)) {
            int currNumber = counter.incrementAndGet();
            if (currNumber >= batchNumber) {
                List<T> tList = new ArrayList<>(batchNumber);
                int i = 0;
                while (i++ < batchNumber) {
                    T outT = batchQueue.remove();
                    tList.add(outT);
                }
                this.batchExecuteHandler.accept(tList);
            }
        }
    }

    public void doAfterComplete() {
        int currentNumber = counter.current();
        if (currentNumber != 0) {
            List<T> list = new ArrayList<>(currentNumber);
            while (!batchQueue.isEmpty()) {
                T outT = batchQueue.poll();
                list.add(outT);
            }
            this.batchExecuteHandler.accept(list);
        }
    }


}
